#!/usr/bin/python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import enum
import glob
import ctypes
import ctypes.util
import random
import shutil
import pathlib
import functools
import subprocess
import atomicwrites


class Util:

    @staticmethod
    def forceDelete(path):
        if os.path.islink(path):
            os.remove(path)
        elif os.path.isfile(path):
            os.remove(path)
        elif os.path.isdir(path):
            shutil.rmtree(path)
        elif os.path.lexists(path):
            os.remove(path)             # other type of file, such as device node
        else:
            pass                        # path does not exist, do nothing

    @staticmethod
    def generateMachineId():
        return ''.join(random.choice('0123456789abcdef') for _ in range(32))

    @staticmethod
    def checkMachineId(machine_id):
        if len(machine_id) != 32:
            return False
        if any([c not in '0123456789abcdef' for c in machine_id]):
            return False
        return True

    @staticmethod
    def isSublistWithOrder(a, b):
        # a: [1, 3, 5]
        # b: [1, 2, 3, 4, 5]
        # returns True

        if any([x not in b for x in a]):
            return False

        c = []
        for x in b:
            if x in a:
                c.append(x)
        if c != a:
            return False

        return True

    @staticmethod
    def compareAndNoneForAll(a, b):
        assert a is not None
        if b is None:
            return True
        else:
            return a == b

    @staticmethod
    def filenameRemoveExt(filename):
        # FIXME
        return filename.replace(".json", "")

    @staticmethod
    def sysDevPathIsPci(devPath):
        return re.search(r"/pci[0-9]{4}:[0-9]{2}/", devPath) is not None

    @staticmethod
    def sysDevPathIsUsb(devPath):
        return re.search(r"/usb[0-9]+/", devPath) is not None

    @classmethod
    def sysDevPathGetMain(cls, target, devPath):
        if target == "sata-target":
            # /sys/devices/pci0000:00/0000:00:01.1/0000:01:00.1/ata6/host5/target5:0:0/5:0:0:0/block/sda -> /sys/devices/pci0000:00/0000:00:01.1/0000:01:00.1/ata6/host5/target5:0:0
            while devPath != "/" and re.fullmatch(r".*/host[0-9]+/target[0-9]+:[0-9]+:[0-9]+", devPath) is None:
                devPath = os.path.dirname(devPath)
        elif target == "pci-device":
            # pci bus can have bridges, go to the first level pci device
            # /sys/devices/pci0000:40/0000:40:01.1/0000:41:00.0/nvme/nvme0/nvme0n1 -> /sys/devices/pci0000:40/0000:40:01.1
            assert cls.sysDevPathIsPci(devPath)
            assert False
        elif target == "pci-device-2":
            # pci bus can have bridges, go to the second level pci device
            # /sys/devices/pci0000:40/0000:40:01.1/0000:41:00.0/nvme/nvme0/nvme0n1 -> /sys/devices/pci0000:40/0000:40:01.1/0000:41:00.0
            assert cls.sysDevPathIsPci(devPath)
            assert False
        elif target == "pci-device-r":
            # pci bus can have bridges, go to the last level pci device
            assert cls.sysDevPathIsPci(devPath)
            while devPath != "/" and re.fullmatch(r"[0-9a-f]{4}:[0-9a-f]{2}:[0-9a-f]+\.[0-9a-f]+", os.path.basename(devPath)) is None:
                devPath = os.path.dirname(devPath)
        elif target == "pci-device-r2":
            # pci bus can have bridges, go to the last but one level pci device
            assert cls.sysDevPathIsPci(devPath)
            assert False
        elif target == "usb-device":
            assert cls.sysDevPathIsUsb(devPath)
            while devPath != "/" and re.fullmatch(r"[0-9]+-[0-9]+(\.[0-9]+)*", os.path.basename(devPath)) is None:
                devPath = os.path.dirname(devPath)
        else:
            assert False
        assert devPath != "/"
        return devPath

    _dmiDecodeCache = dict()

    @classmethod
    def dmiDecodeWithCache(cls, key):
        if key in cls._dmiDecodeCache:
            return cls._dmiDecodeCache[key]

        ret = subprocess.check_output(["dmidecode", "-s", key], universal_newlines=True)
        cls._dmiDecodeCache[key] = ret
        return ret

    @staticmethod
    def udevRuleGetMatchPart(udevRule):
        # FIXME
        assert False

    @staticmethod
    def drmProbe():
        class drmModeRes(ctypes.Structure):
            _fields_ = [
                ("count_fbs", ctypes.c_int),
                ("fbs", ctypes.POINTER(ctypes.c_uint32)),
                ("count_crtcs", ctypes.c_int),
                ("crtcs", ctypes.POINTER(ctypes.c_uint32)),
                ("count_connectors", ctypes.c_int),
                ("connectors", ctypes.POINTER(ctypes.c_uint32)),
                ("count_encoders", ctypes.c_int),
                ("encoders", ctypes.POINTER(ctypes.c_uint32)),
                ("min_width", ctypes.c_uint32),
                ("max_width", ctypes.c_uint32),
                ("min_height", ctypes.c_uint32),
                ("max_height", ctypes.c_uint32),
            ]

        class drmModeConnector(ctypes.Structure):
            _fields_ = [
                ("connector_id", ctypes.c_uint32),
                ("encoder_id", ctypes.c_uint32),
                ("connector_type", ctypes.c_uint32),
                ("connector_type_id", ctypes.c_uint32),
                ("connection", ctypes.c_uint32),
                ("mmWidth", ctypes.c_uint32),
                ("mmHeight", ctypes.c_uint32),
                ("subpixel", ctypes.c_uint32),
                ("count_modes", ctypes.c_int),
                ("modes", ctypes.c_void_p),  # drmModeModeInfo*
                ("count_props", ctypes.c_int),
                ("props", ctypes.POINTER(ctypes.c_uint32)),
                ("prop_values", ctypes.POINTER(ctypes.c_uint64)),
                ("count_encoders", ctypes.c_int),
                ("encoders", ctypes.POINTER(ctypes.c_uint32)),
            ]

        libdrm = ctypes.CDLL(ctypes.util.find_library('drm'))

        libdrm.drmModeGetResources.argtypes = [ctypes.c_int]
        libdrm.drmModeGetResources.restype = ctypes.POINTER(drmModeRes)

        libdrm.drmModeFreeResources.argtypes = [ctypes.POINTER(drmModeRes)]
        libdrm.drmModeFreeResources.restype = None

        libdrm.drmModeGetConnector.argtypes = [ctypes.c_int, ctypes.c_uint32]
        libdrm.drmModeGetConnector.restype = ctypes.POINTER(drmModeConnector)

        libdrm.drmModeFreeConnector.argtypes = [ctypes.POINTER(drmModeConnector)]
        libdrm.drmModeFreeConnector.restype = None

        for fullfn in glob.glob("/dev/dri/card*"):
            fd = os.open(fullfn, os.O_RDWR)
            if fd >= 0:
                try:
                    res = libdrm.drmModeGetResources(fd)
                    if res is not None:
                        try:
                            for i in range(res.contents.count_connectors):
                                drm_conn = libdrm.drmModeGetConnector(fd, res.contents.connectors[i])
                                libdrm.drmModeFreeConnector(drm_conn)
                        finally:
                            libdrm.drmModeFreeResources(res)
                finally:
                    os.close(fd)


class MachineIdFile:

    PATH = "/etc/machine-id"

    @classmethod
    def check(cls):
        buf = cls.read()
        return buf is not None and Util.checkMachineId(buf)

    @classmethod
    def read(cls):
        if os.path.exists(cls.PATH):
            return pathlib.Path(cls.PATH).read_text().strip("\n")
        else:
            return None

    @classmethod
    def generate(cls):
        cls.save(Util.generateMachineId())

    @classmethod
    def save(cls, machine_id):
        if machine_id is None:
            Util.forceDelete(cls.PATH)
        else:
            assert Util.checkMachineId(machine_id)
            with atomicwrites.atomic_write(cls.PATH, overwrite=True) as f:
                f.write(machine_id + "\n")


@functools.total_ordering
class OrderedEnum(enum.Enum):

    @classmethod
    @functools.lru_cache(None)
    def _member_list(cls):
        return list(cls)

    def __lt__(self, other):
        if self.__class__ is other.__class__:
            member_list = self.__class__._member_list()
            return member_list.index(self) < member_list.index(other)
        return NotImplemented


class Expression:

    def __init__(self, expr):
        self._expr = expr

    def evaluate(self, variables):
        nonExistVarNames = []
        for varName, value in variables:
            if value is None:
                nonExistVarNames.append(varName)
            else:
                exec("%s = value" % (varName))
        try:
            return eval(self._expr)
        except NameError as e:
            m = re.fullmatch("name '(.*)' is not defined", str(e))
            if m is not None and m.group(1) in nonExistVarNames:
                return False
            raise
        except ExpressionError:
            return False


class DictForExpression:

    def __init__(self, obj):
        self._obj = obj

    def __getattr__(self, name):
        return getattr(self._obj, name)

    def __getitem__(self, key):
        try:
            return self._obj[key]
        except KeyError as e:
            raise ExpressionError(e)

    def __setitem__(self, key, value):
        self._obj[key] = value

    def __delitem__(self, key):
        del self._obj[key]

    def __len__(self):
        return len(self._obj)

    def __repr__(self):
        return str(self._obj)


class ExpressionError(Exception):

    def __init__(self, e):
        self.original_exception = e
